using System.Collections;
using System.Collections.Generic;
using UnityEngine;

public class PBFObject : MonoBehaviour {
	[SerializeField]
	[Range(0.01f,1.0f)]
	private float radius = default;

	[SerializeField]
	[Range(0.5f, 10f)]
	private float ratio = 1.0f;

	[SerializeField]
	private Mesh mesh = default;
	[SerializeField]
	private Material material = default;

	[SerializeField]
	private ComputeShader solver = default;

	int constraintSolveKernel;
	int XSPHKernel;
	int final1Kernel;
	int final2Kernel;
	int lagrangeKernal;
	int predictorKernel;
	int findNeighborKernel;

	[SerializeField]
	private int particleNum;
	Vector3Int particleBox;

	[System.Serializable]
	struct LiquidParams {
		[Range(0.01f, 100f)]
		public float rho;
		[Range(0f, 100f)]
		public float softConstrainCoef;
		[Range(0f, 0.1f)]
		public float scorrCoef1;
		[Range(1, 10)]
		public int scorrCoef2;
		[Range(0.01f, 10f)]
		public float scorrCoef3;
		[Range(0f, 1f)]
		public float XSPHCoef;
	};
	[SerializeField]
	LiquidParams liquidParams;

	Vector3[] position;
	Vector3[] velocity;
	Bounds bounds;

	[SerializeField]
	private bool showDensity = false;

	ComputeBuffer lambdaBuffer;
	ComputeBuffer positionBuffer;
	ComputeBuffer velocityBuffer;
	ComputeBuffer neighborBuffer;
	ComputeBuffer neighborsCounterBuffer;
	ComputeBuffer velocityXSPHBuffer;
	ComputeBuffer positionPredictBuffer;

	const int threadxNum = 8;
	const int threadyNum = 8;
	const int blockyNum = 64;

	const int maxNeighborNum = 50;

	// Start is called before the first frame update
	void Start() {
		initVariable();
		initparticle();

		material.SetBuffer("positionBuffer", positionBuffer);
		material.SetBuffer("neighborsCounterBuffer", neighborsCounterBuffer);
	}

	void SetKeywords() {
		if (showDensity)
			material.EnableKeyword("_SHOW_DENSITY");
		else
			material.DisableKeyword("_SHOW_DENSITY");
	}
	// Update is called once per frame
	void Update() {
		SetKeywords();

		Graphics.DrawMeshInstancedProcedural(mesh, 0, material, bounds, particleNum);
		//call position change function
		Solve();
	}

	public void OnDrawGizmos() {
		Gizmos.DrawWireCube(this.transform.position, this.transform.localScale);
	}

	private void OnDisable() {
		if (lambdaBuffer != null) lambdaBuffer.Release();
		if (neighborBuffer != null) neighborBuffer.Release();
		if (positionBuffer != null) positionBuffer.Release();
		if (velocityBuffer != null) velocityBuffer.Release();
		if (positionPredictBuffer != null) positionPredictBuffer.Release();
		//positionBuffer = null;
		//velocityBuffer = null;
		//positionPredictBuffer = null;
	}

	private void initVariable() {
		particleBox = new Vector3Int(
			(int)(this.transform.localScale[0] / 2 / radius),
			(int)(this.transform.localScale[1] / 2 / radius),
			(int)(this.transform.localScale[2] / 2 / radius));
		particleNum = particleBox[0] * particleBox[1] * particleBox[2];

		position = new Vector3[particleNum];
		velocity = new Vector3[particleNum];

		bounds = new Bounds(this.transform.position, this.transform.localScale);

		constraintSolveKernel = solver.FindKernel("ConstraintSolve");
		predictorKernel = solver.FindKernel("Predictor");
		findNeighborKernel = solver.FindKernel("FindNeighbor");
		lagrangeKernal = solver.FindKernel("Lagrange");
		final1Kernel = solver.FindKernel("Final1");
		final2Kernel = solver.FindKernel("Final2");
		XSPHKernel = solver.FindKernel("XSPH");

		lambdaBuffer = new ComputeBuffer(particleNum, 4);
		positionBuffer = new ComputeBuffer(particleNum, 3 * 4);
		velocityBuffer = new ComputeBuffer(particleNum, 3 * 4);
		velocityXSPHBuffer = new ComputeBuffer(particleNum, 3 * 4);
		neighborBuffer = new ComputeBuffer(particleNum, maxNeighborNum * 4);
		neighborsCounterBuffer = new ComputeBuffer(particleNum, 4);
		positionPredictBuffer = new ComputeBuffer(particleNum, 3 * 4);
	}
	private void initparticle() {
		Vector3 basepoint = new Vector3(
		   this.transform.position[0] - this.transform.localScale[0] / 2 + radius,
		   this.transform.position[1] - this.transform.localScale[1] / 2 + radius,
		   this.transform.position[2] - this.transform.localScale[2] / 2 + radius);

		for (int i = 0; i < particleBox[0]; i++) {
			int offseti = i * particleBox[1] * particleBox[2];
			for (int j = 0; j < particleBox[1]; j++) {
				int offsetj = j * particleBox[2];
				for (int k = 0; k < particleBox[2]; k++) {
					position[offseti + offsetj + k].Set(
						i * 2 * radius,
						j * 2 * radius,
						k * 2 * radius);
					position[offseti + offsetj + k] += basepoint + new Vector3(Random.Range(-0.1f, 0.1f), Random.Range(-0.1f, 0.1f), Random.Range(-0.1f, 0.1f));
					velocity[offseti + offsetj + k] = Vector3.zero;
				}
			}
		}

		positionBuffer.SetData(position);
		velocityBuffer.SetData(velocity);
		positionPredictBuffer.SetData(position);

		material.SetFloat("scale", radius * 2);
	}

	const int maxSolverIterations = 5;
	private void Solve() {
		//DEBUG
		//positionBuffer.GetData(position);
		//Debug.Log(position[0]);

		//float[] a = new float[10];
		//lambdaBuffer.GetData(a);
		//Debug.Log("lambda : " + a[0]);

		//int[] b = new int[maxNeighborNum+1];
		//neighborBuffer.GetData(b);
		//Debug.Log("neighbor : " + b[0] + " - " + b[1]);


		solver.SetInt("particleNum", particleNum); 
		solver.SetInt("scorrCoef2", liquidParams.scorrCoef2);
		solver.SetFloat("radius", radius);
		solver.SetFloat("h", radius * ratio);
		solver.SetFloat("rho", liquidParams.rho);
		solver.SetFloat("softConstrainCoef", liquidParams.softConstrainCoef);
		solver.SetFloat("scorrCoef1", liquidParams.scorrCoef1);
		solver.SetFloat("scorrCoef3", liquidParams.scorrCoef3);
		solver.SetFloat("XSPHCoef", liquidParams.XSPHCoef);
		solver.SetFloat("xMin", Mathf.Sin(Time.time * 2) * 5 - 5);
		//solver.SetFloat("zMin", 5 - 2 * radius);

		solver.SetBuffer(predictorKernel, "position", positionBuffer);
		solver.SetBuffer(predictorKernel, "velocity", velocityBuffer);
		solver.SetBuffer(predictorKernel, "positionPredict", positionPredictBuffer);
		solver.Dispatch(predictorKernel, 1, 1, 1);
		//solver.Dispatch(kernel, particleNum/blockyNum/threadxNum/threadyNum + 1, blockyNum, 1);

		solver.SetBuffer(findNeighborKernel, "velocity", velocityBuffer);
		solver.SetBuffer(findNeighborKernel, "neighbors", neighborBuffer);
		solver.SetBuffer(findNeighborKernel, "neighborsCounter", neighborsCounterBuffer);
		solver.SetBuffer(findNeighborKernel, "positionPredict", positionPredictBuffer);
		solver.Dispatch(findNeighborKernel, 1, 1, 1);

		for (int i = 0; i < maxSolverIterations; i++) {
			solver.SetBuffer(lagrangeKernal, "lambda", lambdaBuffer);
			solver.SetBuffer(lagrangeKernal, "neighbors", neighborBuffer);
			solver.SetBuffer(lagrangeKernal, "neighborsCounter", neighborsCounterBuffer);
			solver.SetBuffer(lagrangeKernal, "positionPredict", positionPredictBuffer);
			solver.Dispatch(lagrangeKernal, 1, 1, 1);

			solver.SetBuffer(constraintSolveKernel, "lambda", lambdaBuffer);
			solver.SetBuffer(constraintSolveKernel, "position", positionBuffer);
			solver.SetBuffer(constraintSolveKernel, "velocity", velocityBuffer);
			solver.SetBuffer(constraintSolveKernel, "neighbors", neighborBuffer);
			solver.SetBuffer(constraintSolveKernel, "neighborsCounter", neighborsCounterBuffer);
			solver.SetBuffer(constraintSolveKernel, "positionPredict", positionPredictBuffer);
			solver.Dispatch(constraintSolveKernel, 1, 1, 1);
		}

		solver.SetBuffer(final1Kernel, "position", positionBuffer);
		solver.SetBuffer(final1Kernel, "velocity", velocityBuffer);
		solver.SetBuffer(final1Kernel, "positionPredict", positionPredictBuffer);
		solver.Dispatch(final1Kernel, 1, 1, 1);

		solver.SetBuffer(XSPHKernel, "velocity", velocityBuffer);
		solver.SetBuffer(XSPHKernel, "neighbors", neighborBuffer);
		solver.SetBuffer(XSPHKernel, "neighborsCounter", neighborsCounterBuffer);
		solver.SetBuffer(XSPHKernel, "velocityXSPH", velocityXSPHBuffer);
		solver.SetBuffer(XSPHKernel, "positionPredict", positionPredictBuffer);
		solver.Dispatch(XSPHKernel, 1, 1, 1);

		solver.SetBuffer(final2Kernel, "position", positionBuffer);
		solver.SetBuffer(final2Kernel, "velocity", velocityBuffer);
		solver.SetBuffer(final2Kernel, "velocityXSPH", velocityXSPHBuffer);
		solver.SetBuffer(final2Kernel, "positionPredict", positionPredictBuffer);
		solver.Dispatch(final2Kernel, 1, 1, 1);
	}
} 